type Word = u64;
const WORD_BYTES: usize = std::mem::size_of::<Word>();
const WORD_BITS: usize = WORD_BYTES * 8;

pub trait Idx: Copy + std::fmt::Debug + Eq + Ord {
    fn index(self) -> usize;
    fn new(index: usize) -> Self;
    #[inline]
    fn plus(self, offset: usize) -> Self {
        Self::new(self.index() + offset)
    }
    #[inline]
    fn increment_by(&mut self, offset: usize) {
        *self = self.plus(offset);
    }
}

impl Idx for usize {
    fn index(self) -> usize {
        self
    }

    fn new(index: usize) -> Self {
        index
    }
}

impl Idx for u64 {
    fn index(self) -> usize {
        self as usize
    }

    fn new(index: usize) -> Self {
        index as u64
    }
}

/// A growable set of bits.
pub struct BitSet<T> {
    domain_size: usize,
    words: Vec<Word>,
    _marker: std::marker::PhantomData<T>,
}

impl<T> Default for BitSet<T> {
    fn default() -> Self {
        BitSet {
            domain_size: 0,
            words: Vec::new(),
            _marker: std::marker::PhantomData,
        }
    }
}

impl<T: Idx> BitSet<T> {
    /// Gets the domain size.
    pub fn domain_size(&self) -> usize {
        self.domain_size
    }

    /// Ensure that the set can hold at least `min_domain_size` elements.
    pub fn ensure(&mut self, min_domain_size: usize) {
        if self.domain_size < min_domain_size {
            self.domain_size = min_domain_size;
        }

        let min_num_words = num_words(min_domain_size);
        if self.words.len() < min_num_words {
            self.words.resize(min_num_words, 0)
        }
    }

    /// Creates a new, empty bitset.
    #[inline]
    pub fn new() -> Self {
        BitSet::default()
    }

    /// Returns `true` if `self` contains `elem`.
    #[inline]
    pub fn contains(&self, elem: T) -> bool {
        let elem = elem.index();
        elem < self.domain_size && {
            let (word_index, mask) = word_index_and_mask(elem);
            (self.words[word_index] & mask) != 0
        }
    }

    /// Insert `elem`. Returns whether the set has changed.
    #[inline]
    pub fn insert(&mut self, elem: T) -> bool {
        let elem = elem.index();
        self.ensure(elem + 1);
        let (word_index, mask) = word_index_and_mask(elem);
        let word_ref = &mut self.words[word_index];
        let word = *word_ref;
        let new_word = word | mask;
        *word_ref = new_word;
        new_word != word
    }
}

#[inline]
fn num_words(domain_size: usize) -> usize {
    (domain_size + WORD_BITS - 1) / WORD_BITS
}

#[inline]
fn word_index_and_mask<T: Idx>(elem: T) -> (usize, Word) {
    let elem = elem.index();
    let word_index = elem / WORD_BITS;
    let mask = 1 << (elem % WORD_BITS);
    (word_index, mask)
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_bitset() {
        let mut bitset = BitSet::new();
        for i in 0..100_u64 {
            assert!(!bitset.contains(i));
            bitset.insert(i);
            assert!(bitset.contains(i));
            assert!(!bitset.contains(i + 1));
        }
    }
}
